#!/usr/bin/env python3

from outlines import prompt
from .tool_functions import tool_registry


examples = """
    "User": "Hey how are you today?",
    "Assistant": {
        "think_step_by_step": "The user is asking about my well-being. No tools are required to answer this question.",
        "response": "I'm good thanks, how are you?",
        "tool_call": []
    }
    "User": "I'm great, what is the square root of 4?",
    "Assistant": {
        "think_step_by_step": "The user is asking for the square root of 4. I will use the Calculator tool to compute the answer.",
        "response": "Let me calculate that for you.",
        "tool_call": [
            {
                "tool": "sqrt",
                "tool_input": {'a': 4}
            }
        ]
    }
    "Tool": 2.0,
    "Assistant": {
        "think_step_by_step": "After receiving the result from the Calculator tool, I can now inform the user of the result.",
        "response": "It looks like the answer is 2!",
        "tool_call": []
    }
    "User": "Thanks could you tell me what 4 to the power of 2 is?",
    "Assistant": {
        "think_step_by_step": "The user is asking for 4 raised to the power of 2. I will use the Calculator tool to compute the answer.",
        "response": "Let me calculate that for you.",
        "tool_call": [
            {
                "tool": "pow",
                "tool_input": {'a': 4, 'b': 2}
            }
        ]
    }
    "Tool": 16.0,
    "Assistant": {
        "think_step_by_step": "After receiving the result from the Calculator tool, I can now inform the user of the result.",
        "response": "It looks like the answer is 16!",
        "tool_call": []
    }
    "User": "Thanks, could you tell me what's the delay from 192.168.0.1 to 192.168.0.8?",
    "Assistant": {
        "think_step_by_step": "The user is asking for the network delay between two IP addresses. I will use the Ping tool to calculate the delay.",
        "response": "Let me check the delay for you.",
        "tool_call": [
            {
                "tool": "ping",
                "tool_input": {'a': '192.168.0.1', 'b': '192.168.0.8'}
            }
        ]
    }
    "Tool": 153,
    "Assistant": {
        "think_step_by_step": "After receiving the result from the Ping tool, I can now inform the user of the result.",
        "response": "It looks like the delay is 153ms! Not good but usable.",
        "tool_call": []
    }
    User: Can you help me think the performance of my website?

    Assistant:
        {
            "think_step_by_step": "The user is asking for an analysis of their website performance. I will perform three checks: loading speed, uptime status, and SEO score. I will use the Speed Test tool, Uptime Checker tool, and SEO Analyzer tool for these tasks.",
            "response": "Sure! To think the performance of your website, I'll need to check the loading speed, the uptime status, and the SEO score. Let me break this down into subtasks.",
            "tool_call": [
                {
                    "tool": "speed test",
                    "tool_input": {'url': 'https://example.com' }
                },
                {
                    "tool": "uptime checker",
                    "tool_input": {'url': 'https://example.com' }
                },
                {
                    "tool": "seo thinkr",
                    "tool_input": {'url': 'https://example.com' }
                }
            ]
        }

    Tool:

    Speed Test: "Loading speed: 2.5 seconds"
    Uptime Checker: "Uptime status: 99.8%"
    SEO Analyzer: "SEO score: 85/100"
    Assistant:
        {
            "think_step_by_step": "After receiving results from the Speed Test, Uptime Checker, and SEO Analyzer tools, I can now summarize the findings for the user.",
            "response": "Here are the results of the website performance analysis:\n- Loading speed: 2.5 seconds\n- Uptime status: 99.8%\n- SEO score: 85/100\nOverall, your website is performing well, but there might be room for improvement in loading speed.",
            "tool_call": []
        }
"""

rendered_tools = tool_registry.render_text_description()

@prompt
def prompt_template(prompt_with_history=[], rendered_tools=rendered_tools, examples=examples):
    """
    <|begin_of_text|><|start_header_id|>system<|end_header_id|>
    You are an SDN OpenFlow expert assistant that has access to the following set of tools
    within <tools></tools> XML tags.
    Here are the meta messages for each tool:

    <tools>
    {{ rendered_tools }}
    </tools>

    Given the user input, return the name and input of the tool to use.
    Return your response as a JSON blob with 'tool' and 'tool_input' keys.
    You should think about the conversation step by step first, divide the
    task into multiple subtasks, and solve them step by step. Always follow
    the latest language I use.

    Examples
    --------
    {{ examples }}

    Note, the response are output in json, so every quote
    in keys "think_step_by_step" and "response" should be escaped with a backslash
    (\), for example, {"response": "He said, \"Hello!\""}
    Latest Conversation
    --------
    <|eot_id|>

    {{ prompt_with_history }}

    <|eot_id|><|start_header_id|>assistant<|end_header_id|>:
    """

def convert2template(conversation):
    template = ""
    for turn in conversation:
        if turn['role'] == 'user':
            template += f"<|start_header_id|>user<|end_header_id|>\n{turn['content']}<|eot_id|>\n"
        elif turn['role'] == 'assistant':
            template += f"<|start_header_id|>assistant<|end_header_id|>\n{turn['content']}<|eot_id|>\n"
        elif turn['role'] == 'tool':
            template += f"<|start_header_id|>tool<|end_header_id|>\n{turn['content']}<|eot_id|>\n"
    return template
